# coding: utf-8

import numpy as np

class BitWidthConvert:
    
    def __init__(self):
        self.convert_dict = {
            "to_8": self.convert_to_8,
            "to_16": self.convert_to_16,
            "to_24": self.convert_to_24,
            "to_32": self.convert_to_32,
            "to_64": self.convert_to_64
        }
        self.msb_enable = True
    
    def __del__(self):
        pass
    
    def convert_to_8(self, image_data, current_bit_width):
        if current_bit_width == 6:
            image_data_8bit = self.extend_6_to_8(image_data)
            
        return np.array(image_data_8bit, dtype=np.uint8)
        
    def convert_to_16(self, image_data, current_bit_width):
        if current_bit_width == 10:
            image_data_16bit = self.extend_10_to_16(image_data)
        elif current_bit_width == 12:
            image_data_16bit = self.extend_12_to_16(image_data)
        elif current_bit_width == 14:
            image_data_16bit = self.extend_14_to_16(image_data)
        
        return np.array(image_data_16bit, dtype=np.uint8)
    
    def extend_6_to_8(self, image_data):
        """
        | byte0         | byte1         | byte2         | byte3         | 
        |0|1|2|3|4|5|6|7|0|1|2|3|4|5|6|7|0|1|2|3|4|5|6|7|0|1|2|3|4|5|6|7|
        | p1        | p2        | p3        | p4        | ............. |
        """
        
        actual_length = len(image_data)
        image_data_8bit = []
        
        for i in range(0, actual_length, 3):
            image_data_8bit.append( (image_data[i+0] & 0xFC ) >> 2) # 0b11111100
            image_data_8bit.append( ((image_data[i+0] & 0x03 ) << 4) | ( (image_data[i+1] & 0xF0) >> 4)) # 0b00000011, 0b11110000
            image_data_8bit.append( ((image_data[i+1] & 0x0F) << 2) | ( (image_data[i+2] & 0xC0) >> 6)) # 0b00001111, 0b11000000
            image_data_8bit.append( (image_data[i+2] & 0x3F) ) # 0b00111111
            
        return image_data_8bit
        
        
    def extend_10_to_16(self, image_data):
        """
        | byte 0 | byte 1 | byte 2 | byte 3 |     byte 4                    |
        |--------|--------|--------|--------|-------------------------------|
        | p1 h8  | p2 h8  | p3 h8  | p4 h8  | p4 l2 / p3 l2 / p2 l2 / p1 l2 |
        """
        
        actual_length = len(image_data)
        
        image_data_16bit = []
        
        for i in range(0, actual_length, 5):
            for j in range(4):
                image_data_16bit.append((((image_data[i + j] << 2) & 0xFC) | ((image_data[i + 4] >> (j * 2) ) & 0x03 ) )) # low
                image_data_16bit.append(((image_data[i + j] >> 6 ) & 0x03))        # high
                
        return image_data_16bit
    
    def extend_12_to_16(self, image_data):
        """
        | byte0 | byte1 | byte2         | byte3 | byte4 | byte5         | byte6 | byte7 |
        |-------|-------|---------------|-------|-------|---------------|-------|-------|
        | p1 h8 | p2 h8 | p2 l4 / p1 l4 | p3 h8 | p4 h8 | p4 l4 / p3 l4 |.......|.......|
        """
        
        actual_length = len(image_data)
        image_data_16bit = []
        
        for i in range(0, actual_length, 3):
            image_data_16bit.append(image_data[i+0])
            image_data_16bit.append((image_data[i+2] << 4) & 0xF0)
            image_data_16bit.append(image_data[i+1])
            image_data_16bit.append((image_data[i+2] << 0) & 0xF0)
        
        return image_data_16bit
    
    def extend_14_to_16(self, image_data):
        """
        | byte0 | byte1 | byte2 | byte3 |   byte4  |  byte5  |  byte6   | byte7 | byte8 | byty9 |
        |-------|-------|-------|-------|-------------------------------|-------|-------|-------|
        | p1 h8 | p2 h8 | p3 h8 | p4 h8 | p1 l6 / p2 l6 / p3 l6 / p4 l6 | p5 h8 |.......|.......|
        """
        
        actual_length = len(image_data)
        image_data_16bit = []
        
        for i in range(0, actual_length, 7):
            image_data_16bit.append(image_data[i+0])
            image_data_16bit.append((image_data[i+4] & 0xFC) << 0)  # 0b11111100
            
            image_data_16bit.append(image_data[i+1])
            image_data_16bit.append(((image_data[i+4] & 0x03) << 6) | ((image_data[i+5] & 0xF0) >> 2)) # 0b00000011, 0b11110000
            
            image_data_16bit.append(image_data[i+2])
            image_data_16bit.append((image_data[i+5] & 0x0F) << 4 | ((image_data[i+6] & 0xC0) >> 4)) # 0b00001111, 0b11000000
            
            image_data_16bit.append(image_data[i+3])
            image_data_16bit.append((image_data[i+6] & 0x3F) << 2) # 0b00111111
            
        return image_data_16bit
    
    def convert_to_24(self, image_data, current_bit_width):
        pass
    
    def convert_to_32(self, image_data, current_bit_width):
        pass
    
    def convert_to_64(self, image_data, current_bit_width):
        pass
        
    def convert(self, image_data, current_bit_width:int, target_bit_width:int):        
        return self.convert_dict["to_" + str(target_bit_width)](image_data, current_bit_width)

    
if __name__ == '__main__':

    def image_data_generate(number, max_value):
        import random
        random.seed(180710)
        image_data = [random.randint(0, max_value) for i in range(number)]
        return image_data

    bwc = BitWidthConvert()
    
    # --------------------- 6bit ---------------------
    image_data_6bit_1920_1080 = image_data_generate(int(1920*1080*0.75), 255)
    
    for i in range(0, 20, 3):
        print(hex(image_data_6bit_1920_1080[i])[2:], hex(image_data_6bit_1920_1080[i+1])[2:], hex(image_data_6bit_1920_1080[i+2])[2:], bin(image_data_6bit_1920_1080[i])[2:].zfill(8), bin(image_data_6bit_1920_1080[i+1])[2:].zfill(8), bin(image_data_6bit_1920_1080[i+2])[2:].zfill(8))
        
    image_data_8bit_1920_1080 = bwc.convert(image_data_6bit_1920_1080, 6, 8)
    print("image_data_8bit_1920_1080: ", image_data_8bit_1920_1080[:10])
    image_data_8bit_1920_1080.tofile("image_data_6_to_8bit_1920_1080.raw")

    # --------------------- 10bit ---------------------
    image_data_10bit_1920_1080 = image_data_generate(int(1920*1080*1.25), 255)
    
    # for i in range(0, 20, 2):
    #     print(hex(image_data_10bit_1920_1080[i])[2:], hex(image_data_10bit_1920_1080[i+1])[2:], bin(image_data_10bit_1920_1080[i])[2:].zfill(8), bin(image_data_10bit_1920_1080[i+1])[2:].zfill(8))
    
    image_data_16bit_1920_1080 = bwc.convert(image_data_10bit_1920_1080, 10, 16)
    print("image_data_16bit_1920_1080: ", image_data_16bit_1920_1080[:10])
    image_data_16bit_1920_1080.tofile("image_data_10_t0_16bit_1920_1080.raw")
    
    # --------------------- 12bit ---------------------
    image_data_12bit_1920_1080 = image_data_generate(int(1920*1080*1.5), 255)
    
    # for i in range(0, 20, 3):
    #     print(hex(image_data_12bit_1920_1080[i])[2:], hex(image_data_12bit_1920_1080[i+1])[2:], hex(image_data_12bit_1920_1080[i+2])[2:], bin(image_data_12bit_1920_1080[i])[2:].zfill(8), bin(image_data_12bit_1920_1080[i+1])[2:].zfill(8), bin(image_data_12bit_1920_1080[i+2])[2:].zfill(8))
        
    image_data_16bit_1920_1080 = bwc.convert(image_data_12bit_1920_1080, 12, 16)
    print("image_data_16bit_1920_1080: ", image_data_16bit_1920_1080[:10])    
    image_data_16bit_1920_1080.tofile("image_data_12_to_16bit_1920_1080.raw")
    
    # --------------------- 14bit ---------------------
    image_data_14bit_1920_1080 = image_data_generate(int(1920*1080*1.75), 255)
    
    # for i in range(0, 20, 4):
    #     print(hex(image_data_14bit_1920_1080[i])[2:], hex(image_data_14bit_1920_1080[i+1])[2:], hex(image_data_14bit_1920_1080[i+2])[2:], hex(image_data_14bit_1920_1080[i+3])[2:], bin(image_data_14bit_1920_1080[i])[2:].zfill(8), bin(image_data_14bit_1920_1080[i+1])[2:].zfill(8), bin(image_data_14bit_1920_1080[i+2])[2:].zfill(8), bin(image_data_14bit_1920_1080[i+3])[2:].zfill(8))
        
    image_data_16bit_1920_1080 = bwc.convert(image_data_14bit_1920_1080, 14, 16)
    print("image_data_16bit_1920_1080: ", image_data_16bit_1920_1080[:10])
    image_data_16bit_1920_1080.tofile("image_data_14_to_16bit_1920_1080.raw")